from flask import Blueprint, request, jsonify
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import logging
from . import db  # 相对导入
from .models import SentimentAnalysisResult  # 相对导入

# 创建 Blueprint
sentiment_analysis_bp = Blueprint('sentiment_analysis', __name__)

# 配置日志
logging.basicConfig(level=logging.DEBUG)

# 加载预训练的BERT模型和tokenizer
model_path = 'C:/Users/16201/PycharmProjects/four/my_flask_app/app/model/result/checkpoint-final'
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)

# 定义情感分析的路由
@sentiment_analysis_bp.route('/analyze', methods=['POST'])
def analyze_text():
    try:
        # 获取前端传递的文本数据
        data = request.get_json()
        if not data or 'text' not in data:
            raise ValueError("No text provided for analysis")

        text = data['text']
        if not text:
            raise ValueError("Empty text provided for analysis")

        # 调用多次分析情感的函数并取中值
        sentiment_result = analyze_sentiment(text)
        logging.debug("Sentiment analysis result: %s", sentiment_result)

        # 保存预测结果到数据库
        result_entry = SentimentAnalysisResult(
            text=text,
            pos_probability=sentiment_result['pos'],
            neg_probability=sentiment_result['neg'],
            neutral_probability=sentiment_result['neutral']
        )
        db.session.add(result_entry)
        db.session.commit()

        # 生成响应
        response = generate_response(sentiment_result)
        logging.debug("Generated response: %s", response)

        # 返回响应
        return jsonify({'response': response})

    except ValueError as ve:
        logging.error("ValueError occurred: %s", str(ve))
        return jsonify({'error': 'ValueError: ' + str(ve)}), 500

    except KeyError as ke:
        logging.error("KeyError occurred: %s", str(ke))
        return jsonify({'error': 'KeyError: ' + str(ke)}), 500

    except Exception as e:
        logging.error("Error occurred: %s", str(e))
        return jsonify({'error': 'Exception: ' + str(e)}), 500

# 分析情感的函数，多次预测并取中值
def analyze_sentiment(text, num_trials=3):
    try:
        total_pos = 0.0
        total_neg = 0.0
        total_neutral = 0.0

        for _ in range(num_trials):
            # 使用 tokenizer 对文本进行编码
            inputs = tokenizer(text, return_tensors="pt")

            # 使用模型进行预测
            outputs = model(**inputs)

            # 获取 logits
            logits = outputs.logits

            # 计算 softmax 概率
            predictions = torch.softmax(logits, dim=-1)

            # 获取各个情感类别的概率
            pos_prob = predictions[0][1].item()
            neg_prob = predictions[0][0].item()
            neutral_prob = predictions[0][2].item()

            # 累加各次预测的概率
            total_pos += pos_prob
            total_neg += neg_prob
            total_neutral += neutral_prob

        # 计算平均概率
        pos_prob_avg = total_pos / num_trials
        neg_prob_avg = total_neg / num_trials
        neutral_prob_avg = total_neutral / num_trials

        # 构建情感评分字典
        sentiment_score = {'pos': pos_prob_avg, 'neg': neg_prob_avg, 'neutral': neutral_prob_avg}
        return sentiment_score

    except ValueError as ve:
        logging.error("ValueError occurred in analyze_sentiment: %s", str(ve))
        raise ve

    except KeyError as ke:
        logging.error("KeyError occurred in analyze_sentiment: %s", str(ke))
        raise ke

    except Exception as e:
        logging.error("Error occurred in analyze_sentiment: %s", str(e))
        raise e

# 生成响应的函数
def generate_response(sentiment_result):
    # 获取各个情感类别的概率
    pos_prob = sentiment_result['pos']
    neg_prob = sentiment_result['neg']
    neutral_prob = sentiment_result['neutral']

    # 根据概率最大值生成相应的情感描述
    if pos_prob > max(neg_prob, neutral_prob):
        return f"积极情绪概率: {pos_prob:.2f}"
    elif neg_prob > max(pos_prob, neutral_prob):
        return f"消极情绪概率: {neg_prob:.2f}"
    else:
        return f"中性情绪概率: {neutral_prob:.2f}"
